package com.LeeCode;

import java.util.ArrayList;
import java.util.List;

/**
 * 找出数组中的所有 K 近邻下标
 */

public class Code2200 {
    public static void main(String[] args) {
        int[] nums = {2, 2, 2, 2, 2};
        int key = 2, k = 2;
        System.out.println(new Code2200().findKDistantIndices(nums, key, k));
    }

    public List<Integer> findKDistantIndices(int[] nums, int key, int k) {
        int[] diff = new int[nums.length + 1];

        for (int i = 0; i < nums.length; i++) {
            if (nums[i] != key) continue;
            int start = Math.max(0, i - k), end = Math.min(i + k + 1, diff.length - 1);
            diff[start]++;
            diff[end]--;
        }

        List<Integer> ans = new ArrayList<>();
        int sum = 0;
        for (int i = 0; i < nums.length; i++) {
            sum += diff[i];
            if (sum != 0)
                ans.add(i);
        }

        return ans;
    }
}
